from flask import Flask, render_template, request, jsonify, redirect, url_for, flash, send_file, session
from flask_sqlalchemy import SQLAlchemy
from flask_login import LoginManager, UserMixin, login_user, logout_user, login_required, current_user
from werkzeug.security import generate_password_hash, check_password_hash
from werkzeug.utils import secure_filename
import pandas as pd
import networkx as nx
import json
import os
import uuid
from datetime import datetime
import requests
from dotenv import load_dotenv
import io
import base64
import matplotlib

matplotlib.use('Agg')
import matplotlib.pyplot as plt

load_dotenv()

app = Flask(__name__)
app.config['SECRET_KEY'] = 'your-secret-key-here'
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///knowledge_graph.db'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max file size

# 确保上传目录存在
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs('static/graphs', exist_ok=True)

db = SQLAlchemy(app)
login_manager = LoginManager()
login_manager.init_app(app)
login_manager.login_view = 'login'


# 数据库模型
class User(UserMixin, db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(80), unique=True, nullable=False)
    email = db.Column(db.String(120), unique=True, nullable=False)
    password_hash = db.Column(db.String(120), nullable=False)
    is_admin = db.Column(db.Boolean, default=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)

    def set_password(self, password):
        self.password_hash = generate_password_hash(password)

    def check_password(self, password):
        return check_password_hash(self.password_hash, password)


class KnowledgeGraph(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    title = db.Column(db.String(200), nullable=False)
    description = db.Column(db.Text)
    filename = db.Column(db.String(200), nullable=False)
    graph_data = db.Column(db.Text, nullable=False)  # JSON格式的图数据
    is_public = db.Column(db.Boolean, default=False)
    user_id = db.Column(db.Integer, db.ForeignKey('user.id'), nullable=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow)

    user = db.relationship('User', backref=db.backref('graphs', lazy=True))


@login_manager.user_loader
def load_user(user_id):
    return User.query.get(int(user_id))


# AI服务配置
def get_ai_response(prompt, max_tokens=1000):
    """调用AI API获取响应"""
    try:
        # 使用免费的Hugging Face API或其他免费API
        api_url = "https://api-inference.huggingface.co/models/microsoft/DialoGPT-medium"
        headers = {"Authorization": f"Bearer {os.getenv('HUGGINGFACE_API_KEY', 'your-api-key')}"}

        payload = {
            "inputs": prompt,
            "parameters": {
                "max_length": max_tokens,
                "temperature": 0.7
            }
        }

        response = requests.post(api_url, headers=headers, json=payload)
        if response.status_code == 200:
            result = response.json()
            if isinstance(result, list) and len(result) > 0:
                return result[0].get('generated_text', '抱歉，无法获取详细信息。')

        # 备用响应
        return f"关于 '{prompt}' 的详细信息：这是一个重要的知识点，建议进一步学习相关内容。"
    except Exception as e:
        return f"获取详细信息时出现错误：{str(e)}"


def generate_quiz(node_name, node_info=""):
    """生成测试题目"""
    try:
        questions = []
        # 生成10-20道题目的模板
        for i in range(15):
            question = {
                "id": i + 1,
                "question": f"关于{node_name}的问题{i + 1}：以下哪个说法是正确的？",
                "options": [
                    f"{node_name}的特点A",
                    f"{node_name}的特点B",
                    f"{node_name}的特点C",
                    f"{node_name}的特点D"
                ],
                "correct": 0,
                "explanation": f"这是关于{node_name}的详细解析。正确答案是选项A，因为..."
            }
            questions.append(question)

        return questions
    except Exception as e:
        return []


# 路由定义
@app.route('/')
def index():
    if current_user.is_authenticated:
        return redirect(url_for('dashboard'))
    return render_template('index.html')


@app.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form['username']
        email = request.form['email']
        password = request.form['password']

        if User.query.filter_by(username=username).first():
            flash('用户名已存在')
            return render_template('register.html')

        if User.query.filter_by(email=email).first():
            flash('邮箱已被注册')
            return render_template('register.html')

        user = User(username=username, email=email)
        user.set_password(password)
        db.session.add(user)
        db.session.commit()

        flash('注册成功！')
        return redirect(url_for('login'))

    return render_template('register.html')


@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form['username']
        password = request.form['password']
        user = User.query.filter_by(username=username).first()

        if user and user.check_password(password):
            login_user(user)
            return redirect(url_for('dashboard'))
        else:
            flash('用户名或密码错误')

    return render_template('login.html')


@app.route('/logout')
@login_required
def logout():
    logout_user()
    return redirect(url_for('index'))


@app.route('/dashboard')
@login_required
def dashboard():
    # 获取用户的知识图谱
    user_graphs = KnowledgeGraph.query.filter_by(user_id=current_user.id).all()

    # 获取公开的知识图谱
    public_graphs = KnowledgeGraph.query.filter_by(is_public=True).all()

    # 如果是管理员，获取所有图谱
    all_graphs = []
    if current_user.is_admin:
        all_graphs = KnowledgeGraph.query.all()

    return render_template('dashboard.html',
                           user_graphs=user_graphs,
                           public_graphs=public_graphs,
                           all_graphs=all_graphs)


@app.route('/upload', methods=['GET', 'POST'])
@login_required
def upload_file():
    if request.method == 'POST':
        if 'file' not in request.files:
            flash('没有选择文件')
            return redirect(request.url)

        file = request.files['file']
        title = request.form['title']
        description = request.form.get('description', '')
        is_public = 'is_public' in request.form

        if file.filename == '':
            flash('没有选择文件')
            return redirect(request.url)

        if file and allowed_file(file.filename):
            filename = secure_filename(file.filename)
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(filepath)

            # 解析文件并生成图数据
            graph_data = parse_file_to_graph(filepath)

            if graph_data:
                kg = KnowledgeGraph(
                    title=title,
                    description=description,
                    filename=filename,
                    graph_data=json.dumps(graph_data),
                    is_public=is_public,
                    user_id=current_user.id
                )
                db.session.add(kg)
                db.session.commit()

                flash('文件上传成功！')
                return redirect(url_for('view_graph', graph_id=kg.id))
            else:
                flash('文件解析失败')
        else:
            flash('不支持的文件格式')

    return render_template('upload.html')


@app.route('/graph/<int:graph_id>')
@login_required
def view_graph(graph_id):
    graph = KnowledgeGraph.query.get_or_404(graph_id)

    # 权限检查
    if not (graph.user_id == current_user.id or
            graph.is_public or
            current_user.is_admin):
        flash('没有权限查看此图谱')
        return redirect(url_for('dashboard'))

    # 检查是否可以下载
    can_download = (graph.user_id == current_user.id)

    return render_template('graph.html',
                           graph=graph,
                           can_download=can_download)


@app.route('/api/graph/<int:graph_id>/data')
@login_required
def get_graph_data(graph_id):
    graph = KnowledgeGraph.query.get_or_404(graph_id)

    # 权限检查
    if not (graph.user_id == current_user.id or
            graph.is_public or
            current_user.is_admin):
        return jsonify({'error': '没有权限'}), 403

    return jsonify(json.loads(graph.graph_data))


@app.route('/api/node/<int:graph_id>/<node_id>/details')
@login_required
def get_node_details(graph_id, node_id):
    graph = KnowledgeGraph.query.get_or_404(graph_id)

    # 权限检查
    if not (graph.user_id == current_user.id or
            graph.is_public or
            current_user.is_admin):
        return jsonify({'error': '没有权限'}), 403

    graph_data = json.loads(graph.graph_data)
    node = None

    for n in graph_data['nodes']:
        if str(n['id']) == str(node_id):
            node = n
            break

    if not node:
        return jsonify({'error': '节点不存在'}), 404

    # 调用AI获取详细信息
    prompt = f"请详细介绍{node['label']}的相关知识"
    details = get_ai_response(prompt)

    return jsonify({
        'node': node,
        'details': details
    })


@app.route('/api/node/<int:graph_id>/<node_id>/quiz')
@login_required
def get_node_quiz(graph_id, node_id):
    graph = KnowledgeGraph.query.get_or_404(graph_id)

    # 权限检查
    if not (graph.user_id == current_user.id or
            graph.is_public or
            current_user.is_admin):
        return jsonify({'error': '没有权限'}), 403

    graph_data = json.loads(graph.graph_data)
    node = None

    for n in graph_data['nodes']:
        if str(n['id']) == str(node_id):
            node = n
            break

    if not node:
        return jsonify({'error': '节点不存在'}), 404

    # 生成测试题目
    questions = generate_quiz(node['label'])

    return jsonify({
        'node': node,
        'questions': questions
    })


@app.route('/api/graph/<int:graph_id>/download/<format>')
@login_required
def download_graph(graph_id, format):
    graph = KnowledgeGraph.query.get_or_404(graph_id)

    # 只有图谱所有者可以下载
    if graph.user_id != current_user.id:
        return jsonify({'error': '没有权限下载'}), 403

    if format not in ['png', 'jpg']:
        return jsonify({'error': '不支持的格式'}), 400

    # 生成图像文件
    graph_data = json.loads(graph.graph_data)
    image_path = create_graph_image(graph_data, format)

    return send_file(image_path, as_attachment=True)


@app.route('/api/graph/<int:graph_id>/toggle_public', methods=['POST'])
@login_required
def toggle_graph_public(graph_id):
    graph = KnowledgeGraph.query.get_or_404(graph_id)

    if graph.user_id != current_user.id:
        return jsonify({'error': '没有权限'}), 403

    graph.is_public = not graph.is_public
    graph.updated_at = datetime.utcnow()
    db.session.commit()

    return jsonify({'is_public': graph.is_public})


@app.route('/admin/delete_graph/<int:graph_id>', methods=['POST'])
@login_required
def admin_delete_graph(graph_id):
    if not current_user.is_admin:
        return jsonify({'error': '没有权限'}), 403

    graph = KnowledgeGraph.query.get_or_404(graph_id)
    db.session.delete(graph)
    db.session.commit()

    flash('图谱已删除')
    return redirect(url_for('dashboard'))


# 辅助函数
def allowed_file(filename):
    return '.' in filename and \
        filename.rsplit('.', 1)[1].lower() in {'xlsx', 'xls', 'json', 'csv'}


def parse_file_to_graph(filepath):
    """解析文件并生成图数据"""
    try:
        ext = filepath.rsplit('.', 1)[1].lower()

        if ext in ['xlsx', 'xls']:
            df = pd.read_excel(filepath)
        elif ext == 'csv':
            df = pd.read_csv(filepath)
        elif ext == 'json':
            with open(filepath, 'r', encoding='utf-8') as f:
                data = json.load(f)
                if 'nodes' in data and 'edges' in data:
                    return data

        # 从DataFrame创建图数据
        nodes = []
        edges = []

        if len(df.columns) >= 2:
            # 假设前两列是节点关系
            source_col = df.columns[0]
            target_col = df.columns[1]

            node_set = set()

            for idx, row in df.iterrows():
                source = str(row[source_col])
                target = str(row[target_col])

                node_set.add(source)
                node_set.add(target)

                edges.append({
                    'source': source,
                    'target': target,
                    'id': f"edge_{idx}"
                })

            for i, node in enumerate(node_set):
                nodes.append({
                    'id': node,
                    'label': node,
                    'group': i % 5,  # 分组着色
                    'size': 20
                })

        return {
            'nodes': nodes,
            'edges': edges
        }

    except Exception as e:
        print(f"解析文件错误: {e}")
        return None


def create_graph_image(graph_data, format='png'):
    """创建图像文件"""
    G = nx.Graph()

    # 添加节点和边
    for node in graph_data['nodes']:
        G.add_node(node['id'], label=node.get('label', node['id']))

    for edge in graph_data['edges']:
        G.add_edge(edge['source'], edge['target'])

    # 创建图像
    plt.figure(figsize=(12, 8))
    pos = nx.spring_layout(G)
    nx.draw(G, pos, with_labels=True, node_color='lightblue',
            node_size=500, font_size=8, font_weight='bold')

    # 保存图像
    filename = f"graph_{uuid.uuid4().hex}.{format}"
    filepath = os.path.join('static/graphs', filename)
    plt.savefig(filepath, format=format, dpi=300, bbox_inches='tight')
    plt.close()

    return filepath


if __name__ == '__main__':
    with app.app_context():
        db.create_all()

        # 创建管理员用户
        admin = User.query.filter_by(username='admin').first()
        if not admin:
            admin = User(username='admin', email='admin@example.com', is_admin=True)
            admin.set_password('admin123')
            db.session.add(admin)
            db.session.commit()
    app.run(debug=True, port=5000)

